
from librosa.core import audio
import numpy as np
import librosa
import librosa.display
import matplotlib.pyplot as plt
from moviepy.editor import AudioFileClip
import cv2 as cv
from numpy.core.fromnumeric import clip, resize

video_path = '../../../dataset/train/ID1/017.mp4'
audio_path = './test.wav'
img_path = './demo1.jpg'
train_img_path = './resize.jpg'
WIDTH = 40

#   --------------------------------
#       convert video to audio
#   --------------------------------

def video2audio(video_path, audio_path):
    my_audio = AudioFileClip(video_path)
    my_audio.write_audiofile(audio_path)


#   --------------------------------
#       convert audio to image
#   --------------------------------

def audio2img(audio_path, img_path):
    raw_data, raw_framesize = librosa.load(audio_path, sr=None, mono=False)
    print(raw_data.shape)
    print(type(raw_data))
    raw_length = len(raw_data)
    left_data = raw_data[0]
    left_clip_data = raw_data[0, 0:0+100000]
    print(len(left_data))
    print(len(left_clip_data))

    data = left_data * 1.0 / left_data.max()
    framelength = 0.025
    framesize = int(framelength * raw_framesize)

    # get feature
    mel_spect = librosa.feature.melspectrogram(data, sr=raw_framesize, n_fft=framesize)
    mel_spect = librosa.power_to_db(mel_spect, ref=np.max)
    print(type(mel_spect))
    print(mel_spect.shape)

    # draw
    plt.axis('off')
    librosa.display.specshow(mel_spect, sr=raw_framesize, x_axis='time', y_axis='mel')
    plt.savefig(img_path, transparent=False, bbox_inches='tight', pad_inches=0)

    raw_img = cv.imread(img_path)
    print(raw_img.shape)

# #   --------------------------------
# #       clip img with WIDTH
# #   --------------------------------

# raw_img = cv.imread(img_path)
# print(raw_img.shape)
# clip_num = int(raw_img.shape[1] / (WIDTH/2))
# print(clip_num)
# clip_img = []
# for w_idx in range(clip_num-1):
#     clip_img.append(raw_img[:, w_idx*(WIDTH//2):w_idx*(WIDTH//2)+WIDTH, :])
# clip_img.append(raw_img[:, raw_img.shape[1]-WIDTH:raw_img.shape[1], :])
# print(np.array(clip_img).shape)

# #   --------------------------------
# #       resize img to (x,224,224,3)
# #   --------------------------------

# resize_img = []
# for clip_section in clip_img:
#     resize_img.append(cv.resize(clip_section, (224,224), interpolation=cv.INTER_CUBIC))
# cv.imwrite(train_img_path, resize_img[0])
# resize_img = np.array(resize_img)
# print(resize_img.shape)

if __name__ == '__main__':
    video2audio(video_path, audio_path)
    audio2img(audio_path, img_path)
    